#pragma once

#include <ros/ros.h>
#include <ros/console.h>
#include <nav_msgs/Odometry.h>
#include <std_msgs/Float32MultiArray.h>
#include <unitree_legged_msgs/LowState.h>
#include "unitree_legged_sdk/unitree_legged_sdk.h"

#include "convert.h"
#include "interface.hpp"
#include "ekf.hpp"

#include <iostream>
#include <Eigen/Core>
#include <ctime>
#include <cmath>
#include <vector>

using namespace UNITREE_LEGGED_SDK;

#define WINDOW_SIZE 5

class stateEstimator
{
private:
    unitree_legged_msgs::LowState low_state_ros;

    std::string low_state_topic;
    ros::Subscriber low_state_sub;

    std::array<int, 12> joint_idx = {FR_0, FR_1, FR_2,
                                    FL_0, FL_1, FL_2,
                                    RR_0, RR_1, RR_2,
                                    RL_0, RL_1, RL_2};

    ros::Publisher odom_pub;
    ros::Publisher log_pub;
    ros::Publisher filter_pub;

    ros::WallTimer timer;

    int rate = 100;
    float dt = 0.01;
    int16_t threshold;

    std::string child_frame_id, parent_frame_id;

    robotInterface<float> robot_interface;
    extendedKalmanFilter<float> ekf;

    bool verbose = false;
    bool init = false;
    float g_thres;

    int filter_idx = 0;

    Eigen::Matrix<float, WINDOW_SIZE, 4> foot_force_window;
    Eigen::Matrix<float, WINDOW_SIZE, 3> gyroscope_window;
    Eigen::Matrix<float, WINDOW_SIZE, 3> accelerometer_window;
    
public:
    stateEstimator(ros::NodeHandle *nh,
              std::string low_state_topic = "/ros_udp_comm_low/low_state",
              std::string odom_topic = "odom",
              std::string parent_frame = "odom",
              std::string child_frame = "base_ekf",
              int thres = 10)
    {
        auto hints = ros::TransportHints().tcpNoDelay();
        
        low_state_sub = nh->subscribe(low_state_topic, 10, &stateEstimator::lowStateCallback, this, hints);

        odom_pub = nh->advertise<nav_msgs::Odometry>(odom_topic, 10);
        log_pub = nh->advertise<std_msgs::Float32MultiArray>("log", 10);
        filter_pub = nh->advertise<std_msgs::Float32MultiArray>("filter", 10);

        timer = nh->createWallTimer(ros::WallDuration(1.0 / rate), &stateEstimator::periodicControlCallback, this);
        timer.stop();

        robot_interface = robotInterface<float>();
        ekf = extendedKalmanFilter<float>();
        dt = 1.0f / rate;

        child_frame_id = child_frame;
        parent_frame_id = parent_frame;

        threshold = (int16_t)thres;

        foot_force_window.setZero();
        gyroscope_window.setZero();
        accelerometer_window.setZero();

        float process_noise_p_imu = 0.2;
        float process_noise_v_imu = 0.2;
        float process_noise_p_foot = 0.002;
        float sensor_noise_p_imu_rel_foot = 0.005;
        float sensor_noise_v_imu_rel_foot = 0.1;
        float sensor_noise_z_foot = 0.001;
        float high_suspect_num = 1000; 
        nh->getParam("/state_estimation/process_noise_p_imu", process_noise_p_imu);
        nh->getParam("/state_estimation/process_noise_v_imu", process_noise_v_imu);
        nh->getParam("/state_estimation/process_noise_p_foot", process_noise_p_foot);
        nh->getParam("/state_estimation/sensor_noise_p_imu_rel_foot", sensor_noise_p_imu_rel_foot);
        nh->getParam("/state_estimation/sensor_noise_v_imu_rel_foot", sensor_noise_v_imu_rel_foot);
        nh->getParam("/state_estimation/sensor_noise_z_foot", sensor_noise_z_foot);
        nh->getParam("/state_estimation/high_suspect_num", high_suspect_num);
        nh->param<float>("/state_estimation/g_thres", g_thres, 0.1);
        nh->param<bool>("/state_estimation/verbose", verbose, false);
        ekf.setParam(process_noise_p_imu, process_noise_v_imu, process_noise_p_foot, sensor_noise_p_imu_rel_foot, sensor_noise_v_imu_rel_foot, sensor_noise_z_foot, high_suspect_num);
    }

    void lowStateCallback(const unitree_legged_msgs::LowState::ConstPtr &msg)
    {
        low_state_ros = *msg.get();
        init = true;
        updateRobotInterface();
    }

    void updateRobotInterface()
    {
        for (int i = 0; i < 12; ++i)
        {
            robot_interface.joint_pos[i] = low_state_ros.motorState[joint_idx[i]].q;
            robot_interface.joint_vel[i] = low_state_ros.motorState[joint_idx[i]].dq;
        }

        for (int i = 0; i < 4; ++i)
        {
            foot_force_window(filter_idx, i) = low_state_ros.footForce[i];
        }
        for (int i = 0; i < 3; ++i)
        {
            gyroscope_window(filter_idx, i) = low_state_ros.imu.gyroscope[i];
            accelerometer_window(filter_idx, i) = low_state_ros.imu.accelerometer[i];
        }

        robot_interface.body_quat[0] = low_state_ros.imu.quaternion[0];
        robot_interface.body_quat[1] = low_state_ros.imu.quaternion[1];
        robot_interface.body_quat[2] = low_state_ros.imu.quaternion[2];
        robot_interface.body_quat[3] = low_state_ros.imu.quaternion[3];

        // robot_interface.body_ang_vel[0] = low_state_ros.imu.gyroscope[0];
        // robot_interface.body_ang_vel[1] = low_state_ros.imu.gyroscope[1];
        // robot_interface.body_ang_vel[2] = low_state_ros.imu.gyroscope[2];
        
        robot_interface.body_ang_pos[0] = low_state_ros.imu.rpy[0];
        robot_interface.body_ang_pos[1] = low_state_ros.imu.rpy[1];
        robot_interface.body_ang_pos[2] = low_state_ros.imu.rpy[2];

        // robot_interface.foot_force[0] = low_state_ros.footForce[0];
        // robot_interface.foot_force[1] = low_state_ros.footForce[1];
        // robot_interface.foot_force[2] = low_state_ros.footForce[2];
        // robot_interface.foot_force[3] = low_state_ros.footForce[3];

        // robot_interface.body_lin_acc[0] = low_state_ros.imu.accelerometer[0];
        // robot_interface.body_lin_acc[1] = low_state_ros.imu.accelerometer[1];
        // robot_interface.body_lin_acc[2] = low_state_ros.imu.accelerometer[2];

        robot_interface.updateContactStatus(threshold);

        robot_interface.body_ang_vel[0] = gyroscope_window.col(0).mean();
        robot_interface.body_ang_vel[1] = gyroscope_window.col(1).mean();
        robot_interface.body_ang_vel[2] = gyroscope_window.col(2).mean();

        robot_interface.foot_force[0] = foot_force_window.col(0).mean();
        robot_interface.foot_force[1] = foot_force_window.col(1).mean();
        robot_interface.foot_force[2] = foot_force_window.col(2).mean();
        robot_interface.foot_force[3] = foot_force_window.col(3).mean();

        robot_interface.body_lin_acc[0] = accelerometer_window.col(0).mean();
        robot_interface.body_lin_acc[1] = accelerometer_window.col(1).mean();
        robot_interface.body_lin_acc[2] = accelerometer_window.col(2).mean();

        filter_idx ++;
        filter_idx = filter_idx > WINDOW_SIZE - 1 ? filter_idx - WINDOW_SIZE : filter_idx;

        if (verbose)
        {
            std_msgs::Float32MultiArray msg;
            msg.data.resize(20);
            for (int i = 0; i < 4; ++i)
            {
                msg.data.at(i) = low_state_ros.footForce[i];
                msg.data.at(i + 10) = robot_interface.foot_force[i];
            }
            for (int i = 0; i < 3; ++i)
            {
                msg.data.at(i + 4) = low_state_ros.imu.gyroscope[i];
                msg.data.at(i + 7) = low_state_ros.imu.accelerometer[i];
                msg.data.at(i + 14) = robot_interface.body_ang_vel[i];
                msg.data.at(i + 17) = robot_interface.body_lin_acc[i];
            }
            filter_pub.publish(msg);
        }
    }

    void periodicControlCallback(const ros::WallTimerEvent& event)
    {
        // updateRobotInterface();
        if (ekf.is_inited())
        {
            nav_msgs::Odometry msg;
            msg.header.frame_id = parent_frame_id;
            msg.header.stamp = ros::Time::now();
            msg.child_frame_id = child_frame_id;

            ekf.update_estimation(robot_interface, dt);
            msg.pose.pose.position.x = robot_interface.body_lin_pos[0];
            msg.pose.pose.position.y = robot_interface.body_lin_pos[1];
            msg.pose.pose.position.z = robot_interface.body_lin_pos[2];
            msg.pose.pose.orientation.w = robot_interface.body_quat[0];
            msg.pose.pose.orientation.x = robot_interface.body_quat[1];
            msg.pose.pose.orientation.y = robot_interface.body_quat[2];
            msg.pose.pose.orientation.z = robot_interface.body_quat[3];
            msg.twist.twist.linear.x = robot_interface.body_lin_vel[0];
            msg.twist.twist.linear.y = robot_interface.body_lin_vel[1];
            msg.twist.twist.linear.z = robot_interface.body_lin_vel[2];

            odom_pub.publish(msg);

            if (verbose)
            {
                auto xf = ekf.getForward();
                auto xb = ekf.getBackward();
                auto x = ekf.getState();
                auto y = ekf.getMeasure();
                auto yhat = ekf.getPrediction();
                auto error_y = ekf.getError();
                auto pbar = ekf.getStateNoise();
                auto s = ekf.getMeasureNoise();
                auto p = ekf.getP();
                auto q = ekf.getQ();
                auto r = ekf.getR();

                std_msgs::Float32MultiArray msg;
                msg.layout.data_offset = 0;
                msg.layout.dim.resize(2);
                msg.layout.dim[0].label = "channel";
                msg.layout.dim[0].size = 3;
                msg.layout.dim[0].stride = STATE_SIZE * 3;
                msg.layout.dim[1].label = "state";
                msg.layout.dim[1].size = STATE_SIZE;
                msg.layout.dim[1].stride = STATE_SIZE;
                msg.data.resize(STATE_SIZE * 4 + MEAS_SIZE * 4);
                for (int i = 0; i < STATE_SIZE; ++i)
                {
                    msg.data.at(i) = x[i];
                    msg.data.at(i + STATE_SIZE) = xb[i];
                    msg.data.at(i + STATE_SIZE * 2) = xf[i];
                    msg.data.at(i + STATE_SIZE * 3 + MEAS_SIZE * 3) = p[i];
                }
                for (int i = 0; i < MEAS_SIZE; ++i)
                {
                    msg.data.at(i + STATE_SIZE * 3) = y[i];
                    msg.data.at(i + STATE_SIZE * 3 + MEAS_SIZE) = yhat[i];
                    msg.data.at(i + STATE_SIZE * 3 + MEAS_SIZE * 2) = error_y[i];
                    msg.data.at(i + STATE_SIZE * 4 + MEAS_SIZE * 3) = s[i];
                }
                log_pub.publish(msg);
            }
        }
    }

    void start()
    {
        // initialize
        ros::Rate loop_rate(rate);

        while (!init)
        {
            ros::spinOnce();
        }

        ROS_INFO("Calibration will start after 1 seconds");
        for (int i = 0; i < rate; ++i)
        {
            loop_rate.sleep();
            ros::spinOnce();
        }

        ROS_INFO("Calibrating IMU ...");
        std::vector<float> gs;
        float g0 = -1.0f;
        float g1 = -1.0f;
        int count = 0;
        while (count < rate)
        {
            loop_rate.sleep();
            ros::spinOnce();

            g1 = low_state_ros.imu.accelerometer[0] * low_state_ros.imu.accelerometer[0]
                + low_state_ros.imu.accelerometer[1] * low_state_ros.imu.accelerometer[1]
                + low_state_ros.imu.accelerometer[2] * low_state_ros.imu.accelerometer[2];
            gs.push_back(g1);
            if (g0 > 0 && fabs(g1 - g0) > g_thres * 20.0)
            {
                ROS_INFO_STREAM("fabs(g1 - g0): " << fabs(g1 - g0) << "; Retry calibration ...");
                gs.clear();
                g0 = -1.0f;
                count = 0;
                continue;
            }
            g0 = g1;
            ++count;
        }
        g0 = 0.0f;
        for (auto g: gs)
        {
            g0 += g;
        }
        g0 = std::sqrt(g0 / gs.size());
        ekf.set_gravity(g0);
        ROS_INFO_STREAM("Done! g = " << g0 << std::endl);

        for (int i = 0; i < rate * 5; ++i)
        {
            loop_rate.sleep();
            ros::spinOnce();
        }

        updateRobotInterface();
        ekf.init_state(robot_interface);

        timer.start();
        ros::spin();
    }

};